#include "stdafx.h"
#include "../aclib/HexParser.h"
#include "../aclib/BlockData.h"

#include "../aclib/CBC.h"

#include "MCryptor.h"
#include "MLastBlockHandler.h"
#include "MBlockData.h"

namespace AC = AppliedCryptography;

using ::testing::Return;
using ::testing::ReturnRef;
using ::testing::_;
using ::testing::Eq;
using ::testing::Invoke;
using ::testing::Ref;
using ::testing::AtLeast;

template<unsigned int TBlockSize>
void MCryptorProcess(const AC::Block<TBlockSize>& plaintext, AC::Block<TBlockSize>& chyper)
{
	AC::HexParser Ox;
	AC::Block<TBlockSize> key(Ox("404040"));
	chyper = plaintext ^ key;
};

template<unsigned int TKeySize, unsigned int TBlockSize>
AC::Block<TBlockSize> MStealingProcess(const AC::Cryptor<TKeySize,TBlockSize>& cryptor, 
			const AC::Block<TBlockSize>& last, 
			unsigned int lastBlockLength,
			AC::Block<TBlockSize>& secondToLast)
{
	AC::HexParser Ox;
	AC::Block<TBlockSize> key(Ox("800000"));
	secondToLast ^= key;
	return key;
};

TEST(CipherBlockChaining, Encryption)
{
	AC::HexParser Ox;

	MLastBlockEncryption<4, 3> steal;
	MEncryptor<4, 3> cryptor;
	AC::CBC<4, 3, AC::Encryptor<4, 3>, AC::LastBlockEncryption<4, 3>>::Encryption c;

	AC::Block<3> iv(Ox("800000"));
	//AC::BlockData<3> input(Ox("1112132122233132"));
	MBlockData<3> input;
	//AC::BlockData<3> output;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bp1(Ox("111213")), bp2(Ox("212223")), bp3(Ox("313200"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(3));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp1));
	EXPECT_CALL(input, doConstOperatorAccess(1))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp2));
	EXPECT_CALL(input, doConstOperatorAccess(2))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp3));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(2));

	// XORed block passed to encryption
	AC::Block<3> bx1(Ox("911213")), bx2(Ox("F07070"));
	EXPECT_CALL(cryptor, Process(Eq(bx1), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	EXPECT_CALL(cryptor, Process(Eq(bx2), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	/*
	// This is not going to happen as bp3 is handled by mocked ChypertextStealing
	EXPECT_CALL(cryptor, Process(Eq(bp3), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	*/

	// Encrypted blocks passed to output BlockData object
	//AC::Block<3> bc1(Ox("515253")), bc2(Ox("616263")), bc3(Ox("717270"));
	AC::Block<3> bc1(Ox("D15253")), bc2(Ox("B03030"));
	//Block returned by LasBlockHandler
	AC::Block<3> bcc3(Ox("800000"));
	EXPECT_CALL(output, Add(Eq(bc1), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bc2), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bcc3), 2))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bp3), 2, Eq(bc2)))
		.Times(1)
		.WillRepeatedly(Return(bcc3));
	/*
	EXPECT_CALL(steal, Process(_, Eq(bp3), 2, Eq(bc2)))
		.Times(1)
		.WillOnce(Invoke(MStealingProcess<4, 3>))
		.WillOnce(Return(bcc3));
	*/

	c.Process(cryptor, steal, iv, input, output);
};

TEST(CipherBlockChaining, EncryptionOneBlock)
{
	AC::HexParser Ox;

	MLastBlockEncryption<4, 3> steal;
	MEncryptor<4, 3> cryptor;
	AC::CBC<4, 3, AC::Encryptor<4, 3>, AC::LastBlockEncryption<4, 3>>::Encryption c;

	AC::Block<3> iv(Ox("800000"));
	//AC::BlockData<3> input(Ox("1112132122233132"));
	MBlockData<3> input;
	//AC::BlockData<3> output;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bp1(Ox("111213"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(1));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp1));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(3));

	// Encrypted blocks passed to output BlockData object
	AC::Block<3> bc1(Ox("D15253"));
	EXPECT_CALL(output, Add(Eq(iv), 3))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bp1), 3, Eq(iv)))
		.Times(1)
		.WillRepeatedly(Return(bc1));

	c.Process(cryptor, steal, iv, input, output);
};

TEST(CipherBlockChaining, Decryption)
{
	AC::HexParser Ox;

	MLastBlockDecryption<4, 3> steal;
	MDecryptor<4, 3> cryptor;
	AC::CBC<4, 3, AC::Decryptor<4, 3>, AC::LastBlockDecryption<4, 3>>::Decryption c;

	AC::Block<3> iv(Ox("800000"));
	//AC::BlockData<3> input(Ox("1112132122233132"));
	MBlockData<3> input;
	//AC::BlockData<3> output;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bc1(Ox("D15253")), bc2(Ox("B03030"));
	//Block returned by LasBlockHandler
	AC::Block<3> bcc3(Ox("800000"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(3));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bc1));
	EXPECT_CALL(input, doConstOperatorAccess(1))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bc2));
	EXPECT_CALL(input, doConstOperatorAccess(2))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bcc3));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(2));

	EXPECT_CALL(cryptor, Process(Eq(bc1), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	EXPECT_CALL(cryptor, Process(Eq(bc2), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	/*
	// This is not going to happen as bp3 is handled by mocked ChypertextStealing
	EXPECT_CALL(cryptor, Process(Eq(bp3), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	*/

	// Encrypted blocks passed to output BlockData object
	//AC::Block<3> bc1(Ox("515253")), bc2(Ox("616263")), bc3(Ox("717270"));
	AC::Block<3> bp1(Ox("111213")), bp2(Ox("212223")), bp3(Ox("313200"));
	EXPECT_CALL(output, Add(Eq(bp1), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bp2), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bp3), 2))
		.Times(1);

	AC::Block<3> bx2(Ox("F07070"));
	EXPECT_CALL(steal, Process(_, Eq(bcc3), 2, Eq(bx2)))
		.Times(1)
		.WillRepeatedly(Return(bp3));
	/*
	EXPECT_CALL(steal, Process(_, Eq(bp3), 2, Eq(bc2)))
		.Times(1)
		.WillOnce(Invoke(MStealingProcess<4, 3>))
		.WillOnce(Return(bcc3));
	*/

	c.Process(cryptor, steal, iv, input, output);
};

TEST(CipherBlockChaining, DecryptionOneBlock)
{
	AC::HexParser Ox;

	MLastBlockDecryption<4, 3> steal;
	MDecryptor<4, 3> cryptor;
	AC::CBC<4, 3, AC::Decryptor<4, 3>, AC::LastBlockDecryption<4, 3>>::Decryption c;

	AC::Block<3> iv(Ox("800000"));
	MBlockData<3> input;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bc1(Ox("D15253"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(1));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bc1));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(3));

	// Encrypted blocks passed to output BlockData object
	AC::Block<3> bp1(Ox("111213"));
	EXPECT_CALL(output, Add(Eq(iv ^ iv), 3))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bc1), 3, Eq(iv)))
		.Times(1)
		.WillRepeatedly(Return(bp1));

	c.Process(cryptor, steal, iv, input, output);
};

